"""
Statistical script of Youth University Learning
"""
import sys
import os
import csv
import re
import difflib
import datetime
import time
from prettytable import PrettyTable
from colorama import Fore, init


RATIO_THRESHOLD = 0.75                              # 准确率阈值，介于 0 到 1 之间
AVAILABLE_GRADES = ["20", "21", "22"]               # 参加学习的年级范围
AVAILABLE_CLASSES = ["1班", "2班", "3班"]           # 参加学习的班级范围
OPEN_CSV_WITH_EXCEL = False                         # 是否使用 Excel 打开 CSV 文件
ignore_pattern = r'[ 　,:;.，：；。•·་、_\t～~\(\)（）\[\]【】《》→\-—]'
ignore_pattern_1 = r'[ 　,:;.，：；。•·་、_\t～~\(\)（）\[\]【】《》→]'
ignore_pattern_2 = r'[\-—]'


init(autoreset=True)                                # 初始化 colorama


def getClassInfo():
    """
    Get 班级, 班主任, 班级人数 from 班级信息.csv
    """
    class_info_filename = "班级信息.csv"
    # check if the file exists
    if not os.path.isfile(class_info_filename):
        print("班级信息 does not exist")
        sys.exit(1)

    # open the file, which is encoded in UTF-8 with BOM
    with open(class_info_filename, mode="r", encoding='utf-8-sig') as file:
        # read the file
        reader = csv.DictReader(file)

        # create a list to store the data
        classes = []

        # iterate over the file
        for row in reader:
            row_class = []
            # name example: ["20", "级", "学前教育", "（1）班"]
            row_class_name = row["班级"]
            re_class_name_groups = re.search(
                r"(\d{2})(级)(.+?)((（[\d\+]*）)?班)$", row_class_name).groups()
            row_class.append(list(re_class_name_groups[:4]))
            row_class.append(row["班主任"])
            row_class.append(row["班级人数"])
            classes.append(row_class)
    return classes


splitted_classes = getClassInfo()
classes = ["".join(x[0]) for x in splitted_classes]


def isFromSNZX(row, similar_school_table):
    """
    Check if this row is from 山南市职业技术学校
    """
    isFromSNZX = False
    # keywords
    organization = row["单位"]
    organization = re.sub(r'[1壹|①]', '一', organization)
    organization = re.sub(r'[2贰②]', '二', organization)
    must_contains = ["山南"]
    must_not_contains = ["第二", "中等", "二职"]
    either_contains = ["职业", "职校", "职院", "一职"]
    special_alias = ["市职校", "市一职"]
    # check if the organization is 山南市职业技术学校, not 山南市第二中等职业技术学校
    if organization != "":
        if all([x in organization for x in must_contains]):
            if any([x in organization for x in either_contains]):
                re_organization = re.search(
                    r"[(山南)市](.*?)二(.*?)[职校]", organization)
                if not any([x in organization for x in must_not_contains]) and re_organization is None:
                    isFromSNZX = True
                else:
                    similar_school_table.add_row(
                        [row["序号"], row["姓名"], row["单位"]])
        elif any([x == organization for x in special_alias]):
            isFromSNZX = True
    return isFromSNZX


def getMostSimilarClass(str):
    """
    Get most-similar class name from string
    Return class and diff ratio
    """
    if str == "":
        return -1, 0
    else:
        # str replace
        # class
        str = str.replace("3+2", "（3+2）")
        str = str.replace("3+3", "（3+3）")
        str = str.replace("1班", "（1）班")
        str = str.replace("2班", "（2）班")
        str = str.replace("3班", "（3）班")
        # Typos
        str = str.replace("气修", "汽修")
        str = str.replace("学期", "学前")
        str = str.replace("旅遊", "旅游")
        str = str.replace("导游", "旅游")
        str = str.replace("旅馆", "旅游服务与管理")
        # alias
        str = str.replace("20餐", "20级（中餐）烹饪")
        str = str.replace("20级餐", "20级（中餐）烹饪")
        str = str.replace("21中餐", "21级中餐烹饪与营养膳食")
        str = str.replace("21级中餐", "21级中餐烹饪与营养膳食")
        str = str.replace("21级酒店", "21级高星级饭店")
        str = str.replace("21酒店", "21级高星级饭店")
        str = str.replace("21级烹饪", "21级中餐烹饪与营养膳食")
        str = str.replace("21烹饪", "21级中餐烹饪与营养膳食")
        str = str.replace("22级计算机应用", "22级计算机")
        str = str.replace("22计算机应用", "22级计算机")

        # get str grade
        has_grade = False
        str_grade = re.findall(r"\d{2}", str)
        if len(str_grade) > 0 and str_grade[0] in AVAILABLE_GRADES:
            has_grade = True

        # get str class
        has_class = False
        str_class = re.findall(r"（\d）班", str)
        if len(str_class) > 0:
            has_class = True

        # special cases, should be removed when 22 is available
        if not has_grade:
            # 从 21 级开始，「学前」改名为「保育」
            if "学前" in str:
                str = str.replace("学前", "20级学前")
            # 22 级计算机只有一个班，因此如果写了班级，则分配到 21 级
            if "计算机" in str and has_class:
                str = str.replace("计算机", "21级计算机")

        # count ratio
        max_class_index = -1
        max_ratio = 0
        for i in range(len(classes)):
            ratio = difflib.SequenceMatcher(
                None, str, classes[i]).quick_ratio()
            if has_grade and str_grade[0] != splitted_classes[i][0][0]:
                ratio = 0
            if has_class and len(re.findall(r"（\d）班", splitted_classes[i][0][3])) > 0 and str_class[0] != splitted_classes[i][0][3]:
                ratio = 0
            if ratio > max_ratio:
                max_ratio = ratio
                max_class_index = i

        return max_class_index, max_ratio


def getMatchedWords(str_contains_class):
    """
    divide string into matched words and unmatched words
    """
    # 消除信息中用于作为分隔符的字符「+」
    str_contains_class = re.sub(r'[+＋十➕]', ' ', str_contains_class)
    # number
    str_contains_class = re.sub(r'[〇零oO]', '0', str_contains_class)
    str_contains_class = re.sub(r'[一壹|①]', '1', str_contains_class)
    str_contains_class = re.sub(r'[二贰②]', '2', str_contains_class)
    str_contains_class = re.sub(r'[三叁③]', '3', str_contains_class)
    # 填入特殊的 3+2 和 3+3 班级
    str_contains_class = str_contains_class.replace('3 2', '3+2')
    str_contains_class = str_contains_class.replace('3 3', '3+3')
    # 年级简化 2022 -> 22
    str_contains_class = re.sub(r'20(\d{2})', r'\1', str_contains_class)
    # 年级别名
    str_contains_class = str_contains_class.replace('1年级', '22级')
    str_contains_class = str_contains_class.replace('2年级', '21级')
    str_contains_class = str_contains_class.replace('3年级', '20级')
    str_contains_class = str_contains_class.replace('高1', '22级')
    str_contains_class = str_contains_class.replace('高2', '21级')
    str_contains_class = str_contains_class.replace('高3', '20级')
    # 去除分隔符
    str_contains_class = re.sub(ignore_pattern_1, '', str_contains_class)
    # 特殊错别字，即分隔符包含关键信息的情况
    str_contains_class = str_contains_class.replace('-级', '1级')
    str_contains_class = str_contains_class.replace('-班', '1班')
    # 去除分隔符
    str_contains_class = re.sub(ignore_pattern_2, '', str_contains_class)

    splitted_chars = [s for s in str_contains_class]
    splitted_strs = []
    i = 0
    while i < len(splitted_chars):
        if i+1 < len(splitted_chars):
            ii = "".join([splitted_chars[i], splitted_chars[i+1]])
            if ii in AVAILABLE_GRADES:
                splitted_strs.append(ii)
                i += 2
            elif ii in AVAILABLE_CLASSES:
                splitted_strs.append(ii)
                i += 2
            else:
                splitted_strs.append(splitted_chars[i])
                i += 1
        else:
            splitted_strs.append(splitted_chars[i])
            i += 1
    # count ratio
    max_index = 0
    max_ratio = 0
    max_class_index = -1
    max_school_index = 0
    for i in range(len(splitted_strs)):
        first_phrase = "".join(splitted_strs[:i])
        second_phrase = "".join(splitted_strs[i:])
        first_class_index, first_ratio = getMostSimilarClass(first_phrase)
        second_class_index, second_ratio = getMostSimilarClass(
            second_phrase)
        small_phrase = second_phrase if first_ratio > second_ratio else first_phrase
        current_school_index = 0 if first_ratio > second_ratio else 1
        current_class_index = first_class_index if first_ratio > second_ratio else second_class_index
        current_max_ratio = max(first_ratio, second_ratio)
        if any([x in small_phrase for x in AVAILABLE_GRADES]):
            current_max_ratio = 0
        if current_max_ratio > max_ratio:
            max_ratio = current_max_ratio
            max_index = i
            max_class_index = current_class_index
            max_school_index = current_school_index
    first_phrase = "".join(splitted_strs[:max_index])
    second_phrase = "".join(splitted_strs[max_index:])
    ans_another = second_phrase if max_school_index == 0 else first_phrase
    return ans_another, max_class_index, max_ratio


def getFormattedRow(row, low_ratio_table):
    """
    return answer = {
        "班级": "21级XXX（X）班",
        "class_index": 0,
        "学校": "XXXXX",
        "原始数据": "XXXXX"
    }
    """
    answer = {}
    # keywords
    organization = row["单位"]
    name = row["姓名"]

    cleaned_organization = re.sub(ignore_pattern, '', organization)
    cleaned_organization = cleaned_organization.replace('西藏', '')
    cleaned_organization = cleaned_organization.replace('第一', '')
    cleaned_name = re.sub(ignore_pattern, '', name)

    # grade
    grades = ["2", "1", "0", "二一", "二零", "二〇", "二0", "二1",
              "2一", "2零", "2〇", "2|", "二年级", "一年级", "中一", "中二"]
    # find class from name at first
    ratio = -1
    answer_class_index = -1
    answer["提交时间"] = row["提交时间"]

    # special cases, 志愿者 in name
    if "志愿者" in name:
        # 学习前在“姓名栏”输入“志愿者某某某”，名字前一定要记得加“志愿者”三字前缀
        answer["姓名"] = name.replace("志愿者", "").strip()
    else:
        if any([x in name for x in grades]) or len(cleaned_name) > 6:
            # 名字通常最多 6 个字，如果名字超过 6 个字，那么就认为包含班级
            answer["姓名"], answer_class_index, ratio = getMatchedWords(name)
        elif any([x in organization for x in grades]) or len(cleaned_organization) > 9:
            # 单位通常最多 9 个字，如果单位超过 9 个字，那么就认为包含班级
            # example: 山南市职业技术学校
            answer["姓名"] = name
            _, answer_class_index, ratio = getMatchedWords(organization)
        else:
            # TODO: 这种情况下，我们需要从单位中试图找到班级
            answer["姓名"] = name

    answer["序号"] = row["序号"]
    answer["班级"] = classes[answer_class_index] if answer_class_index >= 0 else ""
    answer["class_index"] = answer_class_index if answer_class_index >= 0 else 9999
    answer["学校"] = '山南市职业技术学校'
    answer["原始数据"] = row["姓名"] + ' ' + row["单位"]

    if ratio < RATIO_THRESHOLD and ratio >= 0:
        rounded_ratio = str(round(ratio, 2))
        if ratio < 0.3:
            rounded_ratio = Fore.LIGHTRED_EX + rounded_ratio + Fore.RESET
        elif ratio < 0.5:
            rounded_ratio = Fore.LIGHTMAGENTA_EX + rounded_ratio + Fore.RESET
        # else:
        #     rounded_ratio = Fore.LIGHTYELLOW_EX + rounded_ratio + Fore.RESET
        low_ratio_table.add_row(
            [row["序号"], answer["原始数据"], answer["班级"], rounded_ratio])

    return answer


def openSchoolTable(tips: str, school_path: str):
    print(Fore.LIGHTGREEN_EX + tips)
    if OPEN_CSV_WITH_EXCEL:
        time.sleep(1)
        try:
            os.system(f"start {school_path}")
            print(Fore.LIGHTBLUE_EX + f"\n正在打开「山南市职业技术学校.csv」...")
        except:
            print(Fore.LIGHTYELLOW_EX +
                  "=========================================")
            print(Fore.LIGHTRED_EX +
                  f"未安装可读取 CSV 格式的软件，请手动打开「山南市职业技术学校.csv」")
    else:
        print(Fore.LIGHTBLUE_EX + f"\n请自行打开「山南市职业技术学校.csv」...")
    input()


def main():
    # check if the user has entered the correct number of arguments
    if len(sys.argv) < 2:
        print(Fore.LIGHTRED_EX +
              "使用方法: python statistic.py <csv_file> [cache]")
        sys.exit(1)

    input_csv_file = sys.argv[1]

    # check if the file exists
    if not os.path.isfile(input_csv_file):
        print(Fore.LIGHTRED_EX + "File does not exist")
        sys.exit(1)

    using_cache = False
    if len(sys.argv) > 2 and sys.argv[2] == "cache":
        using_cache = True

    # get today's time 2022-09-14
    today = datetime.datetime.now().strftime("%Y-%m-%d")

    # create temp folder
    work_folder = os.path.dirname(os.path.abspath(input_csv_file))
    temp_folder = os.path.join(work_folder, "temp")
    school_path = os.path.join(temp_folder, f"{today}-山南市职业技术学校.csv")
    if not os.path.exists(temp_folder):
        os.mkdir(temp_folder)

    if not using_cache:
        # open the file, which is encoded in UTF-8 with BOM
        with open(input_csv_file, mode="r", encoding='utf-8-sig') as file:
            # read the file
            reader = csv.DictReader(file)

            # create a list to store the data
            filtered_rows = []
            similar_school_table = PrettyTable(['序号', '姓名', '单位'])

            # iterate over the file
            print(Fore.LIGHTYELLOW_EX +
                  "=========================================\n")
            print(Fore.LIGHTBLUE_EX + "正在筛选包含「山南市职业技术学校」的行...\n")
            for row in reader:
                if isFromSNZX(row, similar_school_table):
                    filtered_rows.append(row)

            similar_school_table.sortby = "单位"
            print(similar_school_table.get_string(title="疑似「山南市职业技术学校」的单位"))

            filtered_rows = sorted(filtered_rows, key=lambda x: x["单位"])
            with open(school_path, mode="w", encoding='utf-8-sig', newline="") as file:
                # create a writer
                writer = csv.DictWriter(
                    file, fieldnames=filtered_rows[0].keys())
                writer.writeheader()
                writer.writerows(filtered_rows)
            print(Fore.LIGHTYELLOW_EX +
                  "\n=========================================\n")

        openSchoolTable("请检查筛选结果，如果没有问题，请按回车键继续...", school_path)

    raw_messages = {}

    while True:
        # init raw messages
        if raw_messages == {}:
            number_set = set()
            with open(school_path, mode="r", encoding='utf-8-sig') as file:
                reader = csv.DictReader(file)
                for row in reader:
                    number_set.add(row["序号"])
            with open(input_csv_file, mode="r", encoding='utf-8-sig') as file:
                # read the file
                reader = csv.DictReader(file)
                # save absolutely raw message
                for d in reader:
                    if d["序号"] in number_set:
                        raw_messages[d["序号"]] = d["姓名"] + ' ' + d["单位"]

        with open(school_path, mode="r", encoding='utf-8-sig') as file:
            reader = csv.DictReader(file)

            # create a list to store the data
            print(Fore.LIGHTYELLOW_EX +
                  "=========================================\n")
            print(Fore.LIGHTBLUE_EX + "正在提取学生信息，并将其匹配到最接近的班级...\n")
            low_ratio_table = PrettyTable(['序号', '原始数据', '匹配班级', '准确率'])
            datas = []
            # iterate over the file
            for row in reader:
                row_keywords = getFormattedRow(row, low_ratio_table)
                datas.append(row_keywords)
            low_ratio_table.sortby = "匹配班级"

            print(low_ratio_table.get_string(
                title=f"准确率低于 {RATIO_THRESHOLD} 的学生信息"))

            print(Fore.LIGHTGREEN_EX + "\n请检查信息，需要编辑请输入 1，不需要编辑请输入 0 或 Enter：\n")
            if input() == "1":
                openSchoolTable("\n编辑完成后，请按回车键继续...\n", school_path)
            else:
                break

    print(Fore.LIGHTYELLOW_EX + "=========================================\n")
    # print the number of students
    print(Fore.LIGHTBLUE_EX + f"统计完毕，本周学习人次为: {len(datas)}。\n")
    print(Fore.LIGHTYELLOW_EX + "=========================================\n\n\n\n")

    datas = sorted(datas, key=lambda x: (x["class_index"], x["提交时间"]))
    count = {}
    for data in datas:
        count[data["class_index"]] = count.get(data["class_index"], 0) + 1

    table1 = [{
        "序号": index + 1,
        "姓名": x["姓名"],
        "班级": x["班级"],
        "学校": x["学校"],
        "提交时间": x["提交时间"],
        "原始数据": raw_messages[x["序号"]],
    } for index, x in enumerate(datas)]

    # write the data to a new csv file
    with open(os.path.join(temp_folder, f"{today}-青年大学习-全校统计.csv"), mode="w", encoding='utf-8-sig', newline="") as file:
        # create a writer
        writer = csv.DictWriter(file, fieldnames=table1[0].keys())

        # write the header
        writer.writeheader()

        # write the data
        writer.writerows(table1)

    table2 = []
    for i in range(len(classes)):
        table2.append({
            "班级": classes[i],
            "班级人数": splitted_classes[i][2],
            "学习人数": count.get(i, 0) if int(count.get(i, 0)) < int(splitted_classes[i][2]) else splitted_classes[i][2],
            "班主任": splitted_classes[i][1],
        })

    # write the data to a new csv file
    with open(os.path.join(temp_folder, f"{today}-青年大学习-各班级明细.csv"), mode="w", encoding='utf-8-sig', newline="") as file:
        # create a writer
        writer = csv.DictWriter(file, fieldnames=table2[0].keys())

        # write the header
        writer.writeheader()

        # write the data
        writer.writerows(table2)


if __name__ == "__main__":
    main()
